import tf from '../tf'
import { LoadStatus } from '../types'
import BaseLoadHandler from './loadHandler/BaseLoadHandler'

/**
 * 卷积模型的父类
 */
export abstract class BaseConvolutionModel {
  public get isLoaded() {
    return this.getLoadHander().getStatus() == LoadStatus.loaded
  }

  /**
   * 需要的输入padding大小 在卷积时需要
   */
  public abstract getInputPadding(): number

  /**
   * 获取加载器
   */
  public abstract getLoadHander(): BaseLoadHandler
  /**
   * 需要的输出padding大小 合成图片块的边缘覆盖需要
   */
  public getOutPadding() {
    //默认不需要边缘覆盖
    return 0
  }

  /**
   * 模型加载
   */
  public async load(): Promise<void> {
    await this.getLoadHander().load()
  }

  /**
   * 初始化输入tensor
   * @param inputShape
   */
  public initInputTensor(input: tf.Tensor3D): tf.Tensor3D {
    let out = input
    if (input.dtype == 'int32') {
      out = tf.tidy(() => input.div(255).asType('float32')) as tf.Tensor3D
      input.dispose()
      input = out
    }
    if (input.shape[2] == 4) {
      out = input.slice(
        [0, 0, 0],
        [input.shape[0], input.shape[1], 3]
      ) as tf.Tensor3D
      input.dispose()
    }
    return out
  }
  /**
   * 转换输出尺寸变化
   * @param width
   * @param height
   * @returns
   */
  public predictSize(width: number, height: number) {
    return { width, height }
  }
  public initResize(width: number, height: number) {
    return { width, height }
  }
  /**
   * 预测
   * @param input
   */
  public abstract predict(input: tf.Tensor3D): tf.Tensor3D
}
